from typing import Any, Dict, Optional, Union, List

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.document import Document
from app.schemas.document import DocumentCreate, DocumentUpdate


async def get_document_by_id(db: AsyncSession, *, id: int) -> Optional[Document]:
    return await db.get(Document, id)


async def get_documents(
    db: AsyncSession, *, skip: int = 0, limit: int = 100,
    project_id: Optional[int] = None, document_type: Optional[str] = None
) -> List[Document]:
    query = select(Document)
    if project_id:
        query = query.where(Document.project_id == project_id)
    if document_type:
        query = query.where(Document.document_type == document_type)
    query = query.offset(skip).limit(limit)
    result = await db.execute(query)
    return result.scalars().all()


async def create_document(db: AsyncSession, *, obj_in: DocumentCreate, creator_id: int) -> Document:
    db_obj = Document(
        title=obj_in.title,
        description=obj_in.description,
        document_type=obj_in.document_type,
        file_path=obj_in.file_path,
        file_size=obj_in.file_size,
        mime_type=obj_in.mime_type,
        version=obj_in.version,
        is_latest=obj_in.is_latest,
        project_id=obj_in.project_id,
        task_id=obj_in.task_id,
        parent_id=obj_in.parent_id,
        creator_id=creator_id,
    )
    db.add(db_obj)
    await db.commit()
    await db.refresh(db_obj)
    return db_obj


async def update_document(
    db: AsyncSession, *, db_obj: Document, obj_in: Union[DocumentUpdate, Dict[str, Any]]
) -> Document:
    if isinstance(obj_in, dict):
        update_data = obj_in
    else:
        update_data = obj_in.dict(exclude_unset=True)
    for field in update_data:
        setattr(db_obj, field, update_data[field])
    db.add(db_obj)
    await db.commit()
    await db.refresh(db_obj)
    return db_obj


async def delete_document(db: AsyncSession, *, id: int) -> Document:
    document = await db.get(Document, id)
    await db.delete(document)
    await db.commit()
    return document